#!/usr/bin/env python3
import sys
import time
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QLineEdit, QPushButton, QMessageBox


def create_folders_in_project(project, folder_names):
    media_pool = project.GetMediaPool()
    if not media_pool:
        print("无法获取媒体池")
        return False

    root_folder = media_pool.GetRootFolder()
    for folder_name in folder_names:
        folder_exists = False
        sub_folders = root_folder.GetSubFolders()
        if isinstance(sub_folders, list):
            for folder in sub_folders:
                if hasattr(folder, 'GetName') and folder.GetName() == folder_name:
                    folder_exists = True
                    break

        if not folder_exists:
            if not media_pool.AddSubFolder(root_folder, folder_name):
                print(f"无法创建文件夹: {folder_name}")
                return False
    return True


def create_projects(num_episodes):
    resolve_path = "/Library/Application Support/Blackmagic Design/DaVinci Resolve/Developer/Scripting/Modules"
    sys.path.append(resolve_path)
    import DaVinciResolveScript as dvr_script

    resolve = dvr_script.scriptapp("Resolve")
    projectManager = resolve.GetProjectManager()

    top_level_folders = ["DI", "IO", "test"]
    for folder in top_level_folders:
        projectManager.CreateFolder(folder)

    projectManager.OpenFolder("DI")
    di_subfolders = ["Feature", "Trailer"]
    for subfolder in di_subfolders:
        projectManager.CreateFolder(subfolder)
        time.sleep(1)

    projectManager.GotoRootFolder()
    time.sleep(1)

    # 处理 "DI/Feature" 文件夹内的工程
    if projectManager.OpenFolder("DI"):
        if projectManager.OpenFolder("Feature"):
            for i in range(1, num_episodes + 1):
                project_name = f"EP{i:02d}"
                new_project = projectManager.CreateProject(project_name)
                if new_project:
                    folder_names = ["SOURCE", "Timelines", "WP", "FROM_VFX", "FLIPCARD", "PTPW", "SOUND", "SUB", "QC"]
                    if not create_folders_in_project(new_project, folder_names):
                        print(f"无法在工程 {project_name} 中创建文件夹层级")
                else:
                    print(f"无法创建工程: {project_name}")
                time.sleep(1)

    projectManager.GotoRootFolder()
    time.sleep(1)

    # 处理 "IO" 文件夹内的工程
    if projectManager.OpenFolder("IO"):
        io_project_names = ["IO_01", "IO_02", "IO_03"]
        io_folder_names = ["SOURCE", "TO_VFX", "QC", "FEATURE"]

        for project_name in io_project_names:
            new_project = projectManager.CreateProject(project_name)
            if new_project:
                if not create_folders_in_project(new_project, io_folder_names):
                    print(f"无法在工程 {project_name} 中创建文件夹层级")
                projectManager.SaveProject()
                projectManager.CloseProject()
            else:
                print(f"无法创建工程: {project_name}")

    QMessageBox.information(None, "完成", "工程文件夹和工程创建完成。")


class App(QWidget):
    def __init__(self):
        super().__init__()
        self.title = '新剧集项目入库插件'
        self.initUI()

    def initUI(self):
        self.setWindowTitle(self.title)
        layout = QVBoxLayout()

        self.label = QLabel('请输入新来的剧集项目共有多少集（纯数字）:', self)
        layout.addWidget(self.label)

        self.textbox = QLineEdit(self)
        layout.addWidget(self.textbox)

        self.button = QPushButton('开始创建', self)
        self.button.clicked.connect(self.on_click)
        layout.addWidget(self.button)

        self.setLayout(layout)
        self.show()

    def on_click(self):
        num_episodes = self.textbox.text()
        try:
            num_episodes = int(num_episodes)
            self.label.setText("正在运行中，请稍后。")
            QApplication.processEvents()  # 更新UI
            create_projects(num_episodes)  # 调用创建工程的函数
        except ValueError:
            QMessageBox.warning(self, '错误', '请输入有效的数字。')


def main():
    app = QApplication(sys.argv)
    ex = App()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()
